package geppetto.reordering;

import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.ExtractedPhrasePair;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.data.ReorderingData;
import geppetto.phraseTable.PhraseTable;
import geppetto.phraseTable.phrase.Phrase;
import gnu.trove.TIntDoubleHashMap;
import gnu.trove.TIntObjectHashMap;

import java.io.PrintStream;
import java.util.HashMap;




public class MSDReorderingCalculator extends AbstractReorderingCalculator {
	
	protected static int MONO_CODE = MSDOrientationCodes.MONO_CODE;
	protected static int SWAP_CODE = MSDOrientationCodes.SWAP_CODE;
	protected static int OTHER_CODE = MSDOrientationCodes.OTHER_CODE;
	protected static int DRIGHT_CODE = MSDOrientationCodes.DRIGHT_CODE;
	protected static int DLEFT_CODE = MSDOrientationCodes.DLEFT_CODE;
	protected static int NOMONO_CODE = MSDOrientationCodes.NOMONO_CODE;
	
	protected static double DEFAULTSMOOTHING = 0.5;	
	
	protected REO_MODEL_TYPE modelType;
	AlignmentsSet set;
	protected TIntDoubleHashMap prevSmoothingFactors = new TIntDoubleHashMap();
	protected TIntDoubleHashMap nextSmoothingFactors = new TIntDoubleHashMap();
		
	public MSDReorderingCalculator(REO_MODEL_TYPE modelType,AlignmentsSet set) {
		this(DEFAULTSMOOTHING,modelType, set);
	}
	
	public MSDReorderingCalculator(double smoothing, REO_MODEL_TYPE modelType,AlignmentsSet set) {
		super();
		this.modelType = modelType;
		this.set = set;
		setSmoothing(smoothing);
	}

	public MSDReorderingCalculator(AlignmentsSet set) {
		this(REO_MODEL_TYPE.REO_MSD, set);
	}
	
	public void setSmoothing(double smoothing){
		prevSmoothingFactors.put(MONO_CODE, smoothing);
		prevSmoothingFactors.put(SWAP_CODE, smoothing);
		prevSmoothingFactors.put(OTHER_CODE, smoothing);
		prevSmoothingFactors.put(DRIGHT_CODE, smoothing);
		prevSmoothingFactors.put(DLEFT_CODE, smoothing);
		prevSmoothingFactors.put(NOMONO_CODE, smoothing);
		nextSmoothingFactors.put(MONO_CODE, smoothing);
		nextSmoothingFactors.put(SWAP_CODE, smoothing);
		nextSmoothingFactors.put(OTHER_CODE, smoothing);
		nextSmoothingFactors.put(DRIGHT_CODE, smoothing);
		nextSmoothingFactors.put(DLEFT_CODE, smoothing);
		nextSmoothingFactors.put(NOMONO_CODE, smoothing);
	}
	
	public ReorderingInfo getReorderingInfo(ExtractedPhrasePair epp){
		 int sentenceNumber = epp.getSentenceNumber();
		 int startS = epp.getSourceStartPos();
		 int endS = epp.getSourceEndPos();
		 int startT = epp.getForeignStartPos();
		 int endT = epp.getForeignEndPos(); 
		 
		Alignment alignment = set.get(sentenceNumber);
		boolean connectedLeftTopP = isAligned(alignment, startS-1, startT-1);
		boolean connectedRightTopP = isAligned(alignment,endS+1, startT-1 );
		boolean connectedLeftTopN = isAligned(alignment,endS+1, endT+1 );
		boolean connectedRightTopN = isAligned(alignment,startS-1, endT+1 ); 
		
		double probLeftP = connectedLeftTopP && !connectedRightTopP ? 1 : 0;
		double probRightP = connectedRightTopP && !connectedLeftTopP ? 1 : 0;
		double probUnkP = probLeftP == 0 && probRightP == 0 ? 1 : 0;
		double probLeftN = connectedLeftTopN && !connectedRightTopN ? 1 : 0;
		double probRightN = connectedRightTopN && !connectedLeftTopN ? 1 : 0;
		double probUnkN = probLeftN == 0 && probRightN == 0 ? 1 : 0; 
		
		ReorderingInfo info = new  ReorderingInfo();
		
		info.addPrevWordOrientation(convertREO_POS(REO_POS.LEFT), probLeftP);
		info.addPrevWordOrientation(convertREO_POS(REO_POS.RIGHT), probRightP);
		info.addPrevWordOrientation(convertREO_POS(REO_POS.UNKNOWN), probUnkP);
		info.addNextWordOrientation(convertREO_POS(REO_POS.LEFT), probLeftN);
		info.addNextWordOrientation(convertREO_POS(REO_POS.RIGHT), probRightN);
		info.addNextWordOrientation(convertREO_POS(REO_POS.UNKNOWN), probUnkN);
		return info;
	}
	
	public String getStringFromCode(int code){
		return MSDOrientationCodes.intToStringHash.get(code);
	}

	public REO_POS getOrientWordModel(boolean connectedLeftTop,
			boolean connectedRightTop, int startF, int endF, int startE,
			int endE) {

		if (connectedLeftTop && !connectedRightTop)
			return REO_POS.LEFT;
		if (modelType == REO_MODEL_TYPE.REO_MONO)
			return REO_POS.UNKNOWN;
		if (!connectedLeftTop && connectedRightTop)
			return REO_POS.RIGHT;
		if (modelType == REO_MODEL_TYPE.REO_MSD)
			return REO_POS.UNKNOWN;
		/*
		for(int indexF=startF-2*unit; (*ge)(indexF, zero) && !connectedLeftTop; indexF=indexF-unit)
		connectedLeftTop = isAligned(sentence, indexF, startE-unit);
		for(int indexF=endF+2*unit; (*lt)(indexF,countF) && !connectedRightTop; indexF=indexF+unit)
		connectedRightTop = isAligned(sentence, indexF, startE-unit);
		if(connectedLeftTop && !connectedRightTop)
		return REO_POS.DRIGHT;
		else if(!connectedLeftTop && connectedRightTop)
		return REO_POS.DLEFT;*/
		return REO_POS.UNKNOWN;
	}
	
	public int convertREO_POS(REO_POS orient){
		switch(orient){
		 	case LEFT: return MONO_CODE;
		 	case RIGHT: return SWAP_CODE;
		 	case DRIGHT: return DRIGHT_CODE;
		 	case DLEFT: return DLEFT_CODE;
		 	case UNKNOWN:
		 	switch(modelType){
		 	case REO_MONO: return NOMONO_CODE;
		 	case REO_MSD: return OTHER_CODE;
		 	case REO_MSLR: return DRIGHT_CODE;
		 	}
		 	break;
		 	} 
		return OTHER_CODE; // this should not happen
	}
	
	protected boolean isAligned(Alignment alignment, int posS, int posT){
		return (alignment._foreignLen == posT && alignment._sourceLen == posS) || 
		(-1 == posT && -1 == posS) || 
		alignment.hasPosition(posS, posT);
	}

	@Deprecated
	public void smoothReorderingInfo(PhraseTable pt, int maxSourceDuration, int maxTargetDuration) {
		ReorderingInfo info = getGlobalReordering(pt,maxSourceDuration, maxTargetDuration);		

		double total = ReorderingInfo.getTotalCount(info);
		System.out.println("smooting totalCount = " + total);
		System.out.print("smoothing: previous - ");
		for(int key : info.getPreviousWordOrientationCounts().keys()){
			double val = info.getPreviousWordOrientationCounts().get(key);
			double newVal = prevSmoothingFactors.get(key) * (val + 0.1) / total;
			System.out.print(getStringFromCode(key) + ":" + prevSmoothingFactors.get(key) + "*(" + val +  "+0.1)/" + total + "=" + newVal + " ");
			prevSmoothingFactors.put(key, newVal);
		}
		System.out.println();
		
		System.out.print("smoothing: following - ");
		for(int key : info.getNextWordOrientationCounts().keys()){
			double val = info.getNextWordOrientationCounts().get(key);
			double newVal = nextSmoothingFactors.get(key) * (val + 0.1) / total;
			System.out.print(getStringFromCode(key) + ":" + nextSmoothingFactors.get(key) + "*(" + val +  "+0.1)/" + total + "=" + newVal + " ");
			nextSmoothingFactors.put(key, newVal);
			}
		System.out.println();
	}
	
	@Deprecated
	public void printMosesReorderingTable(PrintStream out, BilingualCorpus corpus, byte sentenceSource, PhraseTable pt, int maxTargetDuration, int maxSourceDuration){		
		for(int targetLen = 0; targetLen < maxTargetDuration; targetLen++){
			for(int sourceLen = 0; sourceLen < maxSourceDuration; sourceLen++){
				TIntObjectHashMap<TIntObjectHashMap<Phrase>> phrasesT = pt._phraseTable[sourceLen][targetLen];
				int keysT[] = phrasesT.keys();
				for (int i = 0; i < keysT.length; i++) {
					TIntObjectHashMap<Phrase> phrases = phrasesT.get(keysT[i]);
					int[] keys = phrases.keys();
					for (int j = 0; j < keys.length; j++) {
						Phrase p = phrases.get(keys[j]);
						double[] reordPrevValues = getPhrasePrevReorderingFeatures(p.get_reorderingInfo());
						double[] reordNextValues = getPhraseNextReoderingFeatures(p.get_reorderingInfo());						
						
						out.println(p.stringMosesReordering(corpus, pt._sourcePhrasesVocab, pt._targetPhrasesVocab, reordPrevValues, reordNextValues));
					}
				}
			}
		}
	}
	
	@Deprecated
	public double[] getPhrasePrevReorderingFeatures(ReorderingInfo info){
		int[] params = new int[]{MONO_CODE, SWAP_CODE, OTHER_CODE};
		double[] reordPrevValues = new double[params.length];
		
		double total = 0;
		for (int k = 0; k < params.length; k++){
			int param = params[k];
			double smoothing = prevSmoothingFactors.get(param);
			double reordValue = info.getPreviousWordOrientationCounts().get(param);
			reordPrevValues[k] = smoothing + reordValue;
			total += reordPrevValues[k];
		}						
		for (int k = 0; k < params.length; k++){
			reordPrevValues[k] = reordPrevValues[k] / total;
		}
		
		return reordPrevValues;
	}
	
	@Deprecated
	public double[] getPhraseNextReoderingFeatures(ReorderingInfo info){
		int[] params = new int[]{MONO_CODE, SWAP_CODE, OTHER_CODE};
		double[] reordNextValues = new double[params.length];

		double total = 0;
		for (int k = 0; k < params.length; k++){
			int param = params[k];
			double smoothing = nextSmoothingFactors.get(param);
			double reordValue = info.getNextWordOrientationCounts().get(param);
			reordNextValues[k] = smoothing + reordValue;
			total += reordNextValues[k];
		}
		for (int k = 0; k < params.length; k++){
			reordNextValues[k] = reordNextValues[k] / total;
		}
		return reordNextValues;
	}
	
	@Deprecated
	public double[] getPhraseReorderingFeatures(ReorderingInfo info){
		int[] params = new int[]{MONO_CODE, SWAP_CODE, OTHER_CODE};
		double[] reordValues = new double[params.length*2];
		
		double total = 0;
		for (int k = 0; k < params.length; k++){
			int param = params[k];
			double smoothing = prevSmoothingFactors.get(param);
			double reordValue = info.getPreviousWordOrientationCounts().get(param);
			reordValues[k] = smoothing + reordValue;
			total += reordValues[k];
		}						
		for (int k = 0; k < params.length; k++){
			reordValues[k] = reordValues[k] / total;
		}
		
		total = 0;
		for (int k = 0; k < params.length; k++){
			int param = params[k];
			double smoothing = nextSmoothingFactors.get(param);
			double reordValue = info.getNextWordOrientationCounts().get(param);
			reordValues[k+params.length] = smoothing + reordValue;
			total += reordValues[k+params.length];
		}
		for (int k = 0; k < params.length; k++){
			reordValues[k+params.length] = reordValues[k+params.length] / total;
		}
		
		return reordValues;
	}
	
	@Deprecated
	private ReorderingInfo getGlobalReordering(PhraseTable pt,int maxSourceDuration, int maxTargetDuration){
		ReorderingInfo globalReordering = new ReorderingInfo();
		for(int targetLen = 0; targetLen < maxTargetDuration; targetLen++){
			for(int sourceLen = 0; sourceLen < maxSourceDuration; sourceLen++){
				TIntObjectHashMap<TIntObjectHashMap<Phrase>> phrasesT = pt._phraseTable[sourceLen][targetLen];
				int keysT[] = phrasesT.keys();
				for (int i = 0; i < keysT.length; i++) {
					TIntObjectHashMap<Phrase> phrases = phrasesT.get(keysT[i]);
					int[] keys = phrases.keys();
					for (int j = 0; j < keys.length; j++) {
						Phrase p = phrases.get(keys[j]);
						globalReordering.merge(p.get_reorderingInfo());
					}
				}
			}			
		}
		return globalReordering;
	}
}
